# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from dataclasses import dataclass, field, fields
from typing import Any, Dict, List, Literal, Optional, Union

import json

from swift.hub import get_hub
from swift.llm import Processor, Template, get_model_tokenizer, get_template, load_by_unsloth, safe_snapshot_download
from swift.llm.utils import get_ckpt_dir
from swift.plugin import extra_tuners
from swift.utils import (check_json_format, get_dist_setting, get_logger, import_external_file, is_dist, is_master,
                         set_device, use_hf_hub)
from .data_args import DataArguments
from .generation_args import GenerationArguments
from .model_args import ModelArguments
from .quant_args import QuantizeArguments
from .template_args import TemplateArguments

logger = get_logger()


def get_supported_tuners():
    return {'lora', 'full', 'longlora', 'adalora', 'llamapro', 'adapter', 'vera', 'boft', 'fourierft', 'reft', 'bone'
            } | set(extra_tuners.keys())


@dataclass
class CompatArguments:
    #
    ckpt_dir: Optional[str] = None  # 模型的检查点（checkpoint）目录，可选参数，默认为 None
    load_dataset_config: Optional[bool] = None  # 是否加载数据集配置，可选参数，默认为 None
    lora_modules: List[str] = field(default_factory=list)  # LoRA 模块路径列表，默认为空列表

    # 处理 ckpt_dir 参数的内部方法
    def _handle_ckpt_dir(self: 'BaseArguments'):
        # 确保 ckpt_dir 是一个存在的目录，否则抛出断言错误
        assert os.path.isdir(self.ckpt_dir), f'self.ckpt_dir: {self.ckpt_dir}'

        # 检查 ckpt_dir 或其子目录中是否包含特定的配置文件，判断是 adapter 类型还是模型本体
        if (os.path.exists(os.path.join(self.ckpt_dir, 'adapter_config.json'))
                or os.path.exists(os.path.join(self.ckpt_dir, 'default', 'adapter_config.json'))
                or os.path.exists(os.path.join(self.ckpt_dir, 'reft'))):
            # 如果该目录已在 adapters 列表中，直接返回
            if self.ckpt_dir in self.adapters:
                return
            # 否则将 ckpt_dir 插入 adapters 列表头部（优先加载）
            self.adapters.insert(0, self.ckpt_dir)
        else:
            # 如果不是 adapter 配置，则认为是完整模型，赋值给 self.model
            self.model = self.ckpt_dir

        # 清空 ckpt_dir 参数，表示它已经被处理
        self.ckpt_dir = None

        # 打印警告：提示该参数未来版本将被废弃，应改用 --model 或 --adapters 参数
        logger.warning('The `--ckpt_dir` parameter will be removed in `ms-swift>=3.4`. '
                       'Please use `--model`, `--adapters`.')

    # dataclass 初始化完成后执行的钩子方法
    def __post_init__(self: 'BaseArguments'):
        # 如果 ckpt_dir 被设置，调用处理函数
        if self.ckpt_dir is not None:
            self._handle_ckpt_dir()

        # 如果 lora_modules 不为空，则将其加入 adapters 列表（兼容旧参数）
        if len(self.lora_modules) > 0:
            self.adapters += self.lora_modules

            # 打印警告：提示 lora_modules 参数未来版本将废弃，应统一使用 adapters 参数
            logger.warning('The `--lora_modules` parameter will be removed in `ms-swift>=3.4`. '
                           'Please use `--adapters`.')


@dataclass
class BaseArguments(CompatArguments, GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments,
                    ModelArguments):
    """
    BaseArguments 类继承自多个配置类，包含了模型训练、推理、量化、模板、数据处理等参数。

    参数说明:
        tuner_backend (str): 使用的微调器类型，支持 'peft' 或 'unsloth'。
        train_type (str): 训练方式，比如 full、lora 等。
        seed (int): 随机种子，确保结果可复现。
        model_kwargs (Optional[str]): 额外的模型参数，支持 JSON 字符串或字典。
        load_data_args (bool): 是否加载数据集相关配置。
        use_hf (bool): 是否使用 HuggingFace 模型。
        hub_token (Optional[str]): HuggingFace 或 ModelScope 平台的认证 token。
        custom_register_path (List[str]): 注册自定义数据集时使用的 .py 文件路径。
        ignore_args_error (bool): 忽略参数错误，便于 notebook 环境使用。
        use_swift_lora (bool): 是否使用 swift lora。
    """

    # 微调后端选择，支持 'peft' 或 'unsloth'
    tuner_backend: Literal['peft', 'unsloth'] = 'peft'
    # 训练类型，如 full、lora 等
    train_type: str = field(default='lora', metadata={'help': f'train_type choices: {list(get_supported_tuners())}'})
    # 适配器路径列表
    adapters: List[str] = field(default_factory=list)
    # 外部插件路径列表
    external_plugins: List[str] = field(default_factory=list)

    # 随机种子
    seed: int = 42
    # 模型额外参数
    model_kwargs: Optional[Union[dict, str]] = None
    # 是否从 checkpoint 中加载其他参数
    load_args: bool = True
    # 是否从 checkpoint 中加载数据参数
    load_data_args: bool = False

    # 是否使用 HuggingFace 平台
    use_hf: bool = False
    # HuggingFace/ModelScope 授权 token
    hub_token: Optional[str] = field(
        default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'})
    # 自定义数据注册 .py 文件路径
    custom_register_path: List[str] = field(default_factory=list)

    # 是否忽略参数错误，便于 notebook 环境
    ignore_args_error: bool = False
    # 是否使用 swift lora
    use_swift_lora: bool = False

    # 预留训练参数初始化函数（目前未实现）
    def _prepare_training_args(self, training_args: Dict[str, Any]) -> None:
        pass

    # 注册自定义数据集
    def _init_custom_register(self) -> None:
        if isinstance(self.custom_register_path, str):
            self.custom_register_path = [self.custom_register_path]
        if not self.custom_register_path:
            return
        for path in self.custom_register_path:
            import_external_file(path)
        logger.info(f'Successfully registered {self.custom_register_path}.')

    # 导入外部插件
    def _import_external_plugins(self):
        if isinstance(self.external_plugins, str):
            self.external_plugins = [self.external_plugins]
        if not self.external_plugins:
            return
        for external_plugin in self.external_plugins:
            import_external_file(external_plugin)
        logger.info(f'Successfully imported external_plugins: {self.external_plugins}.')

    # 判断是否是 adapter 文件夹
    @staticmethod
    def _check_is_adapter(adapter_dir: str) -> bool:
        if (os.path.exists(os.path.join(adapter_dir, 'adapter_config.json'))
                or os.path.exists(os.path.join(adapter_dir, 'default', 'adapter_config.json'))
                or os.path.exists(os.path.join(adapter_dir, 'reft'))):
            return True
        return False

    # 初始化适配器下载路径
    def _init_adapters(self):
        if isinstance(self.adapters, str):
            self.adapters = [self.adapters]
        self.adapters = [
            safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token) for adapter in self.adapters
        ]

    # 初始化所有参数
    def __post_init__(self):
        # 设置环境变量
        if self.use_hf or use_hf_hub():
            self.use_hf = True
            os.environ['USE_HF'] = '1'

        # 初始化兼容参数 ckpt_dir/lora_modules
        CompatArguments.__post_init__(self)
        # 初始化适配器路径
        self._init_adapters()
        # 初始化 checkpoint 路径
        self._init_ckpt_dir()
        # 注册自定义数据处理
        self._init_custom_register()
        # 导入外部插件
        self._import_external_plugins()
        # 初始化模型额外参数
        self._init_model_kwargs()

        # 获取分布式训练配置
        self.rank, self.local_rank, self.global_world_size, self.local_world_size = get_dist_setting()
        logger.info(f'rank: {self.rank}, local_rank: {self.local_rank}, '
                    f'world_size: {self.global_world_size}, local_world_size: {self.local_world_size}')

        # 如果不是全量训练，校验是否为合法的 adapter
        if self.train_type not in extra_tuners:
            for adapter in self.adapters:
                assert self._check_is_adapter(adapter), (
                    f'`{adapter}` is not an adapter, please try using `--model` to pass it.')

        # 调用所有父类的 __post_init__ 方法
        ModelArguments.__post_init__(self)
        QuantizeArguments.__post_init__(self)
        TemplateArguments.__post_init__(self)
        DataArguments.__post_init__(self)

        # 登录模型平台
        self.hub = get_hub(self.use_hf)
        if self.hub.try_login(self.hub_token):
            logger.info('hub login successful!')

    # 初始化模型额外参数
    def _init_model_kwargs(self):
        self.model_kwargs: Dict[str, Any] = self.parse_to_dict(self.model_kwargs)
        for k, v in self.model_kwargs.items():
            k = k.upper()
            os.environ[k] = str(v)

    # 是否是适配器训练
    @property
    def is_adapter(self) -> bool:
        return self.train_type not in {'full'}

    # 返回支持的微调器类型
    @property
    def supported_tuners(self):
        return get_supported_tuners()

    # 返回可以合并的适配器类型
    @property
    def adapters_can_be_merged(self):
        return {'lora', 'longlora', 'llamapro', 'adalora'}

    # 从 checkpoint 加载参数
    @classmethod
    def from_pretrained(cls, checkpoint_dir: str):
        self = super().__new__(cls)
        self.load_data_args = True
        self.ckpt_dir = checkpoint_dir
        self.load_args_from_ckpt()
        all_keys = list(f.name for f in fields(BaseArguments))
        for key in all_keys:
            if not hasattr(self, key):
                setattr(self, key, None)
        return self

    # 初始化 checkpoint 目录
    def _init_ckpt_dir(self, adapters=None):
        model = self.model or getattr(self, 'mcore_model', None) or getattr(self, 'load', None)
        self.ckpt_dir = get_ckpt_dir(model, adapters or self.adapters)
        if self.ckpt_dir and self.load_args:
            self.load_args_from_ckpt()

    # 从 checkpoint 的 args.json 中加载参数
    def load_args_from_ckpt(self) -> None:
        from ..train_args import TrainArguments
        args_path = os.path.join(self.ckpt_dir, 'args.json')
        assert os.path.exists(args_path), f'args_path: {args_path}'
        with open(args_path, 'r', encoding='utf-8') as f:
            old_args = json.load(f)

        all_keys = list(f.name for f in fields(BaseArguments))
        data_keys = list(f.name for f in fields(DataArguments))

        load_keys = [
            'bnb_4bit_quant_type',
            'bnb_4bit_use_double_quant',
            'train_type',
            'tuner_backend',
            'use_swift_lora',
            'model_name',
            'model_author',
            'split_dataset_ratio',
            'use_chat_template',
        ]

        skip_keys = list(f.name for f in fields(GenerationArguments) + fields(CompatArguments)) + ['adapters']
        if not isinstance(self, TrainArguments):
            skip_keys += ['max_length']

        all_keys = set(all_keys) - set(skip_keys)

        for key, old_value in old_args.items():
            if key not in all_keys or old_value is None:
                continue
            if not self.load_data_args and key in data_keys:
                continue
            value = getattr(self, key, None)
            if value is None or isinstance(value, (list, tuple)) and len(value) == 0 or key in load_keys:
                setattr(self, key, old_value)

        logger.info(f'Successfully loaded {args_path}.')

    # 保存当前参数到 args.json
    def save_args(self, output_dir=None) -> None:
        if is_master():
            output_dir = output_dir or self.output_dir
            os.makedirs(output_dir, exist_ok=True)
            fpath = os.path.join(output_dir, 'args.json')
            logger.info(f'The {self.__class__.__name__} will be saved in: {fpath}')
            with open(fpath, 'w', encoding='utf-8') as f:
                json.dump(check_json_format(self.__dict__), f, ensure_ascii=False, indent=2)

    # 初始化设备（分布式训练时设置）
    def _init_device(self):
        if is_dist():
            set_device()

    # 获取模板对象
    def get_template(self, processor: 'Processor', template_type: Optional[str] = None) -> 'Template':
        template_kwargs = self.get_template_kwargs()
        template_type = template_type or self.template
        template = get_template(template_type, processor, **template_kwargs)
        return template

    # 获取模型处理器（tokenizer + model）
    def get_model_processor(self,
                            *,
                            model=None,
                            model_type=None,
                            model_revision=None,
                            task_type=None,
                            num_labels=None,
                            **kwargs):
        if self.tuner_backend == 'unsloth':
            return load_by_unsloth(self)

        kwargs.update(self.get_model_kwargs())
        kwargs['model_id_or_path'] = model or self.model
        kwargs['model_type'] = model_type or self.model_type
        kwargs['model_revision'] = model_revision or self.model_revision
        kwargs['task_type'] = task_type or self.task_type
        kwargs['num_labels'] = num_labels or self.num_labels

        return get_model_tokenizer(**kwargs)
